import torch
import torch.nn as nn
import torch.optim as optim

# 检查是否可以使用GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 创建一些基本张量
x = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
y = torch.tensor([[2, 4], [6, 8]], dtype=torch.float32)

print("\n基本张量操作:")
print(f"x = \n{x}")
print(f"y = \n{y}")
print(f"x + y = \n{x + y}")
print(f"x * y = \n{x * y}")

# 定义一个简单的神经网络
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(2, 4)
        self.fc2 = nn.Linear(4, 1)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 创建模型实例
model = SimpleNet().to(device)
print("\n模型结构:")
print(model)

# 创建一些示例数据
X = torch.randn(10, 2).to(device)  # 10个样本，每个样本2个特征
y = torch.randn(10, 1).to(device)  # 10个目标值

# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练模型
print("\n开始训练:")
for epoch in range(100):
    # 前向传播
    outputs = model(X)
    loss = criterion(outputs, y)
    
    # 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if (epoch + 1) % 20 == 0:
        print(f'Epoch [{epoch+1}/100], Loss: {loss.item():.4f}')

print("\n训练完成!")

# 测试模型
with torch.no_grad():
    test_input = torch.tensor([[1.0, 2.0]]).to(device)
    prediction = model(test_input)
    print(f"\n测试输入: {test_input}")
    print(f"模型预测: {prediction}") 